package com.yuan.yblog.portal.aop;

import com.yuan.yblog.common.api.ResultCode;
import com.yuan.yblog.common.consts.RedisKey;
import com.yuan.yblog.common.exception.Asserts;
import com.yuan.yblog.common.utils.IpUtils;
import com.yuan.yblog.common.utils.uuid.IdUtils;
import com.yuan.yblog.portal.enums.LimitType;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import com.yuan.yblog.portal.aop.annotation.SlidingWindowLimiter;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;

/**
 * @author hepeiyuan on 2024/5/21 15:50
 * @description 滑动窗口限流切面类
 */
@Slf4j
@Aspect
@Component
public class SlidingWindowLimiterAspect {

    @Autowired
    private RedisTemplate<Object, Object> luaExecuteRedisTemplate;

    @Autowired
    private RedisScript<Long> SlidingWindowLimitScript;

    /**
     * 不同的接口，不同的流量控制
     * map的key为 Limiter.key
     */
    @Around("@annotation(com.yuan.yblog.portal.aop.annotation.SlidingWindowLimiter)")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable{
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        //拿limit的注解
        SlidingWindowLimiter slidingWindowLimiter = method.getAnnotation(SlidingWindowLimiter.class);
        if (slidingWindowLimiter != null) {
            // key作用：不同的接口，不同的流量控制
            String key = slidingWindowLimiter.key();
            String combineKey = getCombineKey(slidingWindowLimiter, joinPoint);
            List<Object> keys = Collections.singletonList(combineKey);

            Long number = luaExecuteRedisTemplate.execute(
                    SlidingWindowLimitScript,
                    keys,
                    slidingWindowLimiter.windowSize(),
                    System.currentTimeMillis(),
                    IdUtils.fastSimpleUUID() + System.currentTimeMillis(),
                    slidingWindowLimiter.windowTimeRange(),
                    RedisKey.LIMIT_KEY_EXPIRE);
            if (number == null || number <= 0) {
                log.info("触发限流，{}s内限制请求数={}, 限流key={}", slidingWindowLimiter.windowTimeRange(), slidingWindowLimiter.windowSize(), key);
                Asserts.fail(ResultCode.REJECT);
            }
        }
        return joinPoint.proceed();
    }

    /**
     * 根据接口，获取一个组合的 key，所谓的组合的 key，就是在注解的 key 属性基础上，再加上方法的完整路径，如果是 IP 模式的话，就再加上 IP 地址。
     * 以 IP 模式为例，最终生成的 key 类似这样：rate_limit:127.0.0.1-com.yuan.yblog.portal.controller.TestController-test。
     * 如果不是 IP 模式，那么生成的 key 中就不包含 IP 地址
     * @param slidingWindowLimiter 限流注解
     * @param point aop切入点
     * @return redis-key
     */
    public String getCombineKey(SlidingWindowLimiter slidingWindowLimiter, JoinPoint point) {
        StringBuilder stringBuilder = new StringBuilder("limit:" + slidingWindowLimiter.key() + ":");
        // 获取当前请求的ip
        if (slidingWindowLimiter.limitType() == LimitType.IP) {
            stringBuilder.append(IpUtils.getIpAddr(((ServletRequestAttributes) RequestContextHolder.currentRequestAttributes()).getRequest())).append(":");
        }
        // 获取接口的包路径
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        Class<?> targetClass = method.getDeclaringClass();
        String className = targetClass.getName();
        // 拼接：[方法路径]-[方法名]
        stringBuilder.append(className.replaceAll("\\.",":")).append(":").append(method.getName());
        return stringBuilder.toString();
    }
}
